package ao.ai.ml.algo.reinforcement.bandit.flat.ucb;

import ao.ai.ml.algo.reinforcement.bandit.model.algo.flat.FlatBandit;
import ao.ai.ml.algo.reinforcement.bandit.model.algo.flat.FlatGambler;
import ao.ai.ml.algo.reinforcement.bandit.model.algo.flat.impl.FlatBanditTrial;
import ao.ai.ml.algo.reinforcement.bandit.model.reward.Reward;

/**
 * User: aostrovsky
 * Date: 10-Feb-2010
 * Time: 7:38:20 PM
 */
public class Ucb1
        implements FlatGambler
{
    //--------------------------------------------------------------------
    private final FlatBandit bandit;
    private final double[]   rewardSum;
    private final int   []   invocations;
    private       int        totalInvocations;


    //--------------------------------------------------------------------
    public Ucb1(FlatBandit machine)
    {
        bandit           = machine;
        rewardSum        = new double[ machine.armCount() ];
        invocations      = new int   [ machine.armCount() ];
        totalInvocations = 0;
    }


    //-------------------------------------------------------------------------
    @Override
    public FlatBandit.Trial play()
    {
        double maxUcb1      = Double.NEGATIVE_INFINITY;
        int    maxUcb1Index = -1;

        for (int i = 0; i < bandit.armCount(); i++)
        {
            if (invocations[i] == 0)
            {
                maxUcb1Index = i;
                break;
            }

            double mean = (rewardSum[i] / invocations[i]);
            double ucb1 = mean + Math.sqrt(
                            (2.0 * Math.log(totalInvocations)) /
                            invocations[i]);
            if (ucb1 > maxUcb1) {
                maxUcb1      = ucb1;
                maxUcb1Index = i;
            }
        }

        Reward reward      = bandit.trial( maxUcb1Index );
        double rewardValue = reward.reward();

        totalInvocations++;
        rewardSum  [ maxUcb1Index ] += rewardValue;
        invocations[ maxUcb1Index ]++;

        return new FlatBanditTrial(
                     reward, maxUcb1Index);
    }
}
